import json
from tasks.bvd.annotation_class import *

# debug
# from annotation_class import *

# class to transfer annotation json to annotation class
class annotation_json_parser:
    def __init__(self, asm_json_name=None, decomplie_json_name=None):
        self.asm_file_name = asm_json_name
        self.decomplie_file_name = decomplie_json_name

    def get_sensitive_asm_addrs(self, annotation):
        # 获得敏感函数牵扯到所有的asm地址列表，也就是注解中caller_func的basic_block地址字段，内容为元组，表示起始地址和终止地址
        sensitive_asm_addrs = []
        if annotation.plt_sensitive_addr_dict != {}:
            for (func_name, plt_sensitive_func) in annotation.plt_sensitive_addr_dict.items():
                sensitive_asm_addrs.append((plt_sensitive_func.func_mem, plt_sensitive_func.func_mem))
        if annotation.direct_sensitive_dict != {}:
            for (func_name, direct_sensitive_func) in annotation.direct_sensitive_dict.items():
                direct_call_list = direct_sensitive_func.caller_func_info_ls
                for direct_call in direct_call_list:
                    bb_list = direct_call.basic_block
                    for bb in bb_list:
                        sensitive_asm_addrs.append((bb["Block Start"], bb["Block End"]))
        if annotation.inlined_func_dict != {}:
            for (func_name, inline_sensitive_func) in annotation.inlined_func_dict.items():
                inline_call_list = inline_sensitive_func.caller_func_info_ls
                for inline_func in inline_call_list:
                    bb_list = inline_func.basic_block
                    for bb in bb_list:
                        sensitive_asm_addrs.append((bb["Block Start"], bb["Block End"]))
        # print(sensitive_asm_addrs)
        return sensitive_asm_addrs
    
    def parser_plt(self, plt_sensitive_addr_dict):
        plt_dict = {}
        for key in plt_sensitive_addr_dict.keys():
            plt_func = plt_sensitive_func()
            plt_func.func_name = key
            plt_func.func_mem = plt_sensitive_addr_dict[key]['mem']
            plt_func.func_str = plt_sensitive_addr_dict[key]['str']
            plt_dict[key] = plt_func
        return plt_dict

    def parser_direct(self, direct_sensitive_ls):
        direct_dict = {}
        for direct_sensitive in direct_sensitive_ls:
            if direct_sensitive['name'] not in direct_dict.keys():
                
                # direct func info init
                direct_func = direct_sensitive_func()
                direct_func.func_name = direct_sensitive['name']
                # direct_func.inline = direct_sensitive['inline']
                direct_func.inline = False
                direct_func.caller_func_info_ls = []
                direct_dict[direct_sensitive['name']] = direct_func

                # caller func info init
                caller_func_info = asm_caller_func()
                caller_func_info.func_name = direct_sensitive['caller_func']['caller_func_name']
                caller_func_info.call_address = direct_sensitive['call_address']
                caller_func_info.code = direct_sensitive['code']
                caller_func_info.basic_block = direct_sensitive['caller_func']['basic_block']
                caller_func_info.related_par_address = direct_sensitive['related_par_address']
                direct_func.caller_func_info_ls.append(caller_func_info)
            
            else:

                # caller func info init
                caller_func_info = asm_caller_func()
                caller_func_info.func_name = direct_sensitive['caller_func']['caller_func_name']
                caller_func_info.call_address = direct_sensitive['call_address']
                caller_func_info.code = direct_sensitive['code']
                caller_func_info.basic_block = direct_sensitive['caller_func']['basic_block']
                caller_func_info.related_par_address = direct_sensitive['related_par_address']
                direct_dict[direct_sensitive['name']].caller_func_info_ls.append(caller_func_info)

        return direct_dict

    def parser_inline(self, inline_func_ls):
        inline_dict = {}
        for inline_func in inline_func_ls:
            if inline_func['inline_func_name_predict'] not in inline_dict.keys():
                
                # inline func info init
                inline_func_info = inline_sensitive_func()
                inline_func_info.func_name_predict = inline_func['inline_func_name_predict']
                inline_func_info.inline = True
                inline_func_info.caller_func_info_ls = []
                inline_dict[inline_func['inline_func_name_predict']] = inline_func_info

                # caller func info init
                caller_func_info = asm_caller_func()
                caller_func_info.func_name = inline_func['caller_func']['name']
                caller_func_info.call_address = inline_func['call_address']
                caller_func_info.code = inline_func['code']
                caller_func_info.basic_block = inline_func['caller_func']['basic_block']
                inline_dict[inline_func['inline_func_name_predict']].caller_func_info_ls.append(caller_func_info)
            
            else:

                # caller func info init
                caller_func_info = asm_caller_func()
                caller_func_info.func_name = inline_func['caller_func']['name']
                caller_func_info.call_address = inline_func['call_address']
                caller_func_info.code = inline_func['code']
                caller_func_info.basic_block = inline_func['caller_func']['basic_block']
                inline_dict[inline_func['inline_func_name_predict']].caller_func_info_ls.append(caller_func_info)

        return inline_dict

    def parser_decompile(self, decompile_sensitive_ls):
        decompile_dict = {}
        for decompile_sensitive in decompile_sensitive_ls:
            if (decompile_sensitive['name'], decompile_sensitive['inline']) not in decompile_dict.keys():
                decompile_func = decomplie_sensitive_func()
                decompile_func.func_name = decompile_sensitive['name']
                if decompile_sensitive['inline']:
                    decompile_func.inline = True
                else:
                    decompile_func.inline = False
                decompile_func.caller_func_info_ls = []
                decompile_dict[(decompile_sensitive['name'], decompile_sensitive['inline'])] = decompile_func

                # caller func info init
                caller_func_info = decomplie_caller_func()
                caller_func_info.func_name = decompile_sensitive['caller_func_name']
                caller_func_info.caller_func_list = decompile_sensitive['caller_func_list']
                caller_func_info.code = decompile_sensitive['code']
                caller_func_info.call_address = decompile_sensitive['call_address']
                caller_func_info.related_par_ls = decompile_sensitive['related_par']
                decompile_dict[(decompile_sensitive['name'], decompile_sensitive['inline'])].caller_func_info_ls.append(caller_func_info)

            else:
                # caller func info init
                caller_func_info = decomplie_caller_func()
                caller_func_info.func_name = decompile_sensitive['caller_func_name']
                caller_func_info.caller_func_list = decompile_sensitive['caller_func_list']
                caller_func_info.code = decompile_sensitive['code']
                caller_func_info.call_address = decompile_sensitive['call_address']
                caller_func_info.related_par_ls = decompile_sensitive['related_par']
                decompile_dict[(decompile_sensitive['name'], decompile_sensitive['inline'])].caller_func_info_ls.append(caller_func_info)
        
        return decompile_dict

    def construct(self): # naive version
        if self.asm_file_name is not None:
            with open(self.asm_file_name, 'r') as asm_file:
                asm_data = json.load(asm_file)
        else:
            print('[ERROR] No asm annotation found')
        if self.decomplie_file_name is not None:
            with open(self.decomplie_file_name, 'r') as f:
                decompile_data = json.load(f)
        else:
            print('[ERROR] No decompile annotation found')
        ret = annotation(self.asm_file_name, self.decomplie_file_name)
        for key in asm_data.keys():
            if key == 'plt_sensitive_address':
                if asm_data[key] != {}:
                    ret.plt_sensitive_addr_dict = self.parser_plt(asm_data[key])
                else:
                    ret.plt_sensitive_addr_dict = {}
            elif key == 'direct_sensitive_func':
                if asm_data[key] != []:
                    ret.direct_sensitive_dict = self.parser_direct(asm_data[key])
                else:
                    ret.direct_sensitive_dict = {}
            elif key == 'inline_sensitive_func':
                if asm_data[key] != []:
                    ret.inlined_func_dict = self.parser_inline(asm_data[key])
                else:
                    ret.inlined_func_dict = {}
        for key in decompile_data.keys():
            if key == 'sensitive_func':
                if decompile_data[key] != []:
                    ret.decomplie_sensitive_dict = self.parser_decompile(decompile_data[key])
                else:
                    ret.decomplie_sensitive_dict = {}
        return ret


if __name__  == '__main__':
    parser = annotation_json_parser(r'/disk1/zghhhhhhhhhhhhhhhhhhhh/zz_static/zz_static_project/preprocessed_data/notation/CWE121_Stack_Based_Buffer_Overflow_O0/s07/bad/CWE121_Stack_Based_Buffer_Overflow__CWE806_char_declare_snprintf_01/result.json', r'/disk1/zghhhhhhhhhhhhhhhhhhhh/zz_static/zz_static_project/preprocessed_data/notation/CWE121_Stack_Based_Buffer_Overflow_O0/s07/bad/CWE121_Stack_Based_Buffer_Overflow__CWE806_char_declare_snprintf_01/decompile_result.json')
    annotation = parser.construct()
    parser.get_sensitive_asm_addrs(annotation)
    # print(annotation.decomplie_sensitive_dict)